#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
на вход параметры:
N - число вершин в одном измерении в сетке
B - чила магнитного поля
'''
from __future__ import division
import numpy

#так как областью интереса является двумерная сетка, то логично нумеровать вершины как x и y, каждая переменная
#пробегает от 0 до N - 1


def i2xy(i, N):
    '''
    Функция по номеру состояния восстанавливает координаты на сетке
    '''
    return i % N, int(i / N)

def i2state(i):
    '''
    функция по данному i генерирует массив из 1 и -1, что соответствует значиниям спина в соответствующих точках
    '''
    binary = int(str(i), 2)
    binary_array = [ int(j) for j in str(binary) ]
    state = [ 2 * j - 1 for j in binary_array ]
    return state

def state2i(state):
    '''
    функция по заданному массиву из +1 и -1, говорит состоянию с каким номером данный массив соответствует
    '''
    binary_array = [ int((j+1)/2) for j in state ]
    binary = int( ''.join(binary_array), 10 )
    return binary

class Hamiltonian:
    '''
    Класс описывает наш гамильтониан
    на вход хочет параметры:
    neighbors - tuple описывающий элементарную ячейку
    N - число вершин в одном измерении. Работаем с двумерной сеткой, следовательно общее число вершин N x N
    по горизонтали нумеруем x, по вертикали - y 
    ''' 
    def __init__(self, neighbors, N ):
        self.N = N
        #нужно сгерерировать словарь, в котором будут указаны соседи у каждой вершины в формате xy 
        self.xy_neighbors = {}
        for x in xrange(N):
            for y in xrange(N):
                th = []
                for neighbor in neighbors:
                    tn = numpy.array( neighbor) + (x, y)
                    for i in xrange(len(tn) - 1):
                        if tn[i] >= self.N:
                            tn[i] = 0
                    th += [ tuple(tn) ]
                self.xy_neighbors[ (x, y) ] = th
        
        def xy2i(x, y):
            '''
            Функция на вход получает координаты вершины в сетке, а на выход номер вершины в нумерации, где вершине 
            соотвествует один номер. 
            '''
            return y * self.N + x
        
        #Теперь переведём этот словарь в формат, когда каждой вершине сопостовляется одно число
        self.i_neighbors = {}
        for x in xrange(N):
            for y in xrange(N):
                 self.i_neighbors[ xy2i( x, y ) ] = [ xy2i (tx,ty) for tx,ty in self.xy_neighbors[ (x,y) ] ]
                 
        
    def sum_neighbors(self, site, state):
        '''
        функция вычисляет сумму соседей спина site в данном состоянии state 
        '''    
        result = 0 
        for i in self.i_neighbors[ site ]:
            result += state[i]
        return result
    
    def matrix_gen(self, B):
        '''
        Функция генерирует нам матрицу, рписывающую шамильтониан
        '''
        H = numpy.zeros(2 * [N] )
        nStates = 4**N
        for i in xrange(nStates):
            state = i2state(i)
            H_J = 0
            for site in state:
                H_J -= site * self.sum_neighbors(site, state)
        
        H[(i,i)] = H_J

        for k in xrange(4**N):
            state[k] *= -1
            j = state2i( state )
            state[k] *= -1
            H[(i,j)] = B
        return H
    
if __name__ == '__main__':
    N = 2
    B = 0
    neighbors = ( (0, 1), (0, -1), (1, 0), (-1, 0), (-1, -1), (1, 1))
    Hm = Hamiltonian(neighbors, N)
    print Hm.matrix_gen(B)